#include "ransac_ground_filter.hpp"

// #include <pcl_ros/transforms.hpp>  // This header may not exist in some ROS versions
#include <pcl/common/centroid.h>
#include <pcl/point_types.h>

#include <limits>
#include <memory>
#include <random>
#include <string>
#include <vector>

namespace
{
Eigen::Vector3d getArbitraryOrthogonalVector(const Eigen::Vector3d & input)
{
  const double x = input.x();
  const double y = input.y();
  const double z = input.z();
  const double x2 = std::pow(x, 2);
  const double y2 = std::pow(y, 2);
  const double z2 = std::pow(z, 2);

  Eigen::Vector3d unit_vec{0, 0, 0};
  if (x2 <= y2 && x2 <= z2) {
    unit_vec.x() = 0;
    unit_vec.y() = z;
    unit_vec.z() = -y;
    unit_vec = unit_vec / std::sqrt(y2 + z2);
  } else if (y2 <= x2 && y2 <= z2) {
    unit_vec.x() = -z;
    unit_vec.y() = 0;
    unit_vec.z() = x;
    unit_vec = unit_vec / std::sqrt(z2 + x2);
  } else if (z2 <= x2 && z2 <= y2) {
    unit_vec.x() = y;
    unit_vec.y() = -x;
    unit_vec.z() = 0;
    unit_vec = unit_vec / std::sqrt(x2 + y2);
  }
  return unit_vec;
}

geometry_msgs::Pose getDebugPose(const Eigen::Affine3d & plane_affine)
{
  geometry_msgs::Pose debug_pose;
  const Eigen::Quaterniond quat(plane_affine.rotation());
  debug_pose.position.x = plane_affine.translation().x();
  debug_pose.position.y = plane_affine.translation().y();
  debug_pose.position.z = plane_affine.translation().z();
  debug_pose.orientation.x = quat.x();
  debug_pose.orientation.y = quat.y();
  debug_pose.orientation.z = quat.z();
  debug_pose.orientation.w = quat.w();
  return debug_pose;
}
}  // namespace

namespace autoware_ground_segmentation
{

PlaneBasis getPlaneBasis(const Eigen::Vector3d & plane_normal)
{
  PlaneBasis basis;
  basis.e_z = plane_normal;
  basis.e_x = getArbitraryOrthogonalVector(plane_normal);
  basis.e_y = basis.e_x.cross(basis.e_z);
  return basis;
}

RANSACGroundFilter::RANSACGroundFilter(ros::NodeHandle& nh, ros::NodeHandle& private_nh)
: nh_(nh), private_nh_(private_nh)
{
  // Get parameters
  private_nh_.param<std::string>("unit_axis", unit_axis_, "z");
  private_nh_.param<int>("max_iterations", max_iterations_, 1000);
  private_nh_.param<int>("min_inliers", min_inliers_, 100);
  private_nh_.param<int>("min_points", min_points_, 10);
  private_nh_.param<double>("outlier_threshold", outlier_threshold_, 0.1);
  private_nh_.param<double>("plane_slope_threshold", plane_slope_threshold_, 10.0);
  private_nh_.param<double>("height_threshold", height_threshold_, 0.1);
  private_nh_.param<double>("voxel_size_x", voxel_size_x_, 0.1);
  private_nh_.param<double>("voxel_size_y", voxel_size_y_, 0.1);
  private_nh_.param<double>("voxel_size_z", voxel_size_z_, 0.1);
  private_nh_.param<bool>("debug", debug_, false);

  // Set unit vector based on axis
  if (unit_axis_ == "x") {
    unit_vec_ = Eigen::Vector3d::UnitX();
  } else if (unit_axis_ == "y") {
    unit_vec_ = Eigen::Vector3d::UnitY();
  } else {
    unit_vec_ = Eigen::Vector3d::UnitZ();
  }

  // Setup PCL filters
  voxel_grid_.setLeafSize(voxel_size_x_, voxel_size_y_, voxel_size_z_);
  
  seg_.setOptimizeCoefficients(true);
  seg_.setModelType(pcl::SACMODEL_PLANE);
  seg_.setMethodType(pcl::SAC_RANSAC);
  seg_.setDistanceThreshold(outlier_threshold_);
  seg_.setMaxIterations(max_iterations_);

  // Setup ROS
  pointcloud_sub_ = nh_.subscribe("input", 1, &RANSACGroundFilter::pointcloudCallback, this);
  ground_cloud_pub_ = nh_.advertise<sensor_msgs::PointCloud2>("output/ground", 1);
  no_ground_cloud_pub_ = nh_.advertise<sensor_msgs::PointCloud2>("output/no_ground", 1);
  
  if (debug_) {
    debug_pose_array_pub_ = nh_.advertise<geometry_msgs::PoseArray>("debug/plane_poses", 1);
    debug_ground_cloud_pub_ = nh_.advertise<sensor_msgs::PointCloud2>("debug/ground_cloud", 1);
  }

  ROS_INFO("RANSAC Ground Filter initialized");
}

void RANSACGroundFilter::pointcloudCallback(const sensor_msgs::PointCloud2::ConstPtr& input_msg)
{
  // Convert to PCL
  pcl::PointCloud<PointType>::Ptr input_cloud(new pcl::PointCloud<PointType>);
  pcl::fromROSMsg(*input_msg, *input_cloud);

  if (input_cloud->empty()) {
    ROS_WARN("Empty input point cloud");
    return;
  }

  // Voxel grid filtering
  pcl::PointCloud<PointType>::Ptr voxel_cloud(new pcl::PointCloud<PointType>);
  voxel_grid_.setInputCloud(input_cloud);
  voxel_grid_.filter(*voxel_cloud);

  if (voxel_cloud->size() < static_cast<size_t>(min_points_)) {
    ROS_WARN("Not enough points after voxel filtering");
    return;
  }

  // Apply RANSAC
  pcl::PointIndices::Ptr ground_inliers(new pcl::PointIndices);
  pcl::ModelCoefficients::Ptr ground_coefficients(new pcl::ModelCoefficients);
  
  applyRANSAC(voxel_cloud, ground_inliers, ground_coefficients);

  if (ground_inliers->indices.empty()) {
    ROS_WARN("No ground plane found");
    return;
  }

  // 添加调试信息：打印RANSAC选择的点的坐标
  ROS_INFO("=== RANSAC Point Selection Debug ===");
  ROS_INFO("Total voxel cloud points: %zu", voxel_cloud->size());
  ROS_INFO("RANSAC selected inliers: %zu", ground_inliers->indices.size());
  ROS_INFO("First 10 selected points coordinates:");
  for (size_t i = 0; i < std::min(static_cast<size_t>(100), ground_inliers->indices.size()); ++i) {
    const auto& point = voxel_cloud->points[ground_inliers->indices[i]];
    ROS_INFO("  Point %zu: [%.3f, %.3f, %.3f]", i, point.x, point.y, point.z);
  }
  ROS_INFO("================================");

  // Extract ground and non-ground points
  pcl::PointCloud<PointType>::Ptr ground_cloud(new pcl::PointCloud<PointType>);
  pcl::PointCloud<PointType>::Ptr no_ground_cloud(new pcl::PointCloud<PointType>);
  
  extractPointsIndices(voxel_cloud, *ground_inliers, ground_cloud, no_ground_cloud);

  // Check plane slope
  if (ground_coefficients->values.size() >= 4) {
    Eigen::Vector3d plane_normal(
      ground_coefficients->values[0],
      ground_coefficients->values[1],
      ground_coefficients->values[2]
    );
    plane_normal.normalize();
    
    // Debug: Print plane normal and unit vector values
    ROS_INFO("=== DEBUG: Plane Analysis ===");
    ROS_INFO("Raw coefficients: [%.6f, %.6f, %.6f, %.6f]", 
             ground_coefficients->values[0], ground_coefficients->values[1], 
             ground_coefficients->values[2], ground_coefficients->values[3]);
    
    // 验证平面方程的正确性
    ROS_INFO("Plane equation: %.6fx + %.6fy + %.6fz + %.6f = 0", 
             ground_coefficients->values[0], ground_coefficients->values[1], 
             ground_coefficients->values[2], ground_coefficients->values[3]);
    
    // 计算原始法向量的长度
    double raw_normal_length = std::sqrt(
      std::pow(ground_coefficients->values[0], 2) + 
      std::pow(ground_coefficients->values[1], 2) + 
      std::pow(ground_coefficients->values[2], 2));
    ROS_INFO("Raw normal vector length: %.6f", raw_normal_length);
    
    ROS_INFO("Normalized plane_normal: [%.6f, %.6f, %.6f]", 
             plane_normal.x(), plane_normal.y(), plane_normal.z());
    ROS_INFO("Unit vector (unit_vec_): [%.6f, %.6f, %.6f]", 
             unit_vec_.x(), unit_vec_.y(), unit_vec_.z());
    
    // 修复平面斜率计算逻辑
    // 对于地面检测，我们关心的是平面法向量与重力方向（Z轴）的夹角
    // 当平面水平时，法向量应该与Z轴平行，夹角为0度
    // 当平面垂直时，法向量应该与Z轴垂直，夹角为90度
    const auto plane_slope = std::acos(std::abs(plane_normal.dot(unit_vec_))) * 180.0 / M_PI;
    double dot_product = plane_normal.dot(unit_vec_);
    
    ROS_INFO("Dot product: %.6f", dot_product);
    ROS_INFO("Absolute dot product: %.6f", std::abs(dot_product));
    ROS_INFO("Plane slope: %.2f degrees", plane_slope);
    
    // 添加额外的平面验证逻辑
    // 检查平面高度分布，地面平面应该有较小的Z轴变化
    double min_z = std::numeric_limits<double>::max();
    double max_z = std::numeric_limits<double>::lowest();
    for (const auto& point : ground_cloud->points) {
      min_z = std::min(min_z, static_cast<double>(point.z));
      max_z = std::max(max_z, static_cast<double>(point.z));
    }
    double z_range = max_z - min_z;
    ROS_INFO("Ground plane Z range: %.6f (min: %.6f, max: %.6f)", z_range, min_z, max_z);
    
    // 如果Z轴变化太大，可能是垂直表面
    if (z_range > 0.5) {  // 0.5米的阈值
      ROS_WARN("Ground plane has large Z variation (%.3f m), likely a vertical surface", z_range);
      return;
    }
    
    ROS_INFO("================================");
    
    if (plane_slope > plane_slope_threshold_) {
      ROS_WARN("Ground plane slope too steep: %.2f degrees", plane_slope);
      ROS_WARN("This usually means RANSAC detected a vertical surface (wall) instead of ground");
      ROS_WARN("Consider adjusting RANSAC parameters or checking input data quality");
      return;
    }
  }

  // 修改为与ROS2版本一致的高度过滤方法
  // 获取地面平面变换矩阵
  Eigen::Affine3d plane_affine = getPlaneAffine(*ground_cloud, 
    Eigen::Vector3d(ground_coefficients->values[0], 
                    ground_coefficients->values[1], 
                    ground_coefficients->values[2]));

  // 使用高度阈值过滤非地面点
  pcl::PointCloud<PointType>::Ptr filtered_no_ground_cloud(new pcl::PointCloud<PointType>);
  for (const auto& point : *input_cloud) {
    // 将点转换到地面平面坐标系
    const Eigen::Vector3d transformed_point = plane_affine.inverse() * Eigen::Vector3d(point.x, point.y, point.z);
    // 如果点的Z坐标绝对值大于高度阈值，则认为是非地面点
    if (std::abs(transformed_point.z()) > height_threshold_) {
      filtered_no_ground_cloud->points.push_back(point);
    }
  }

  // Publish results
  sensor_msgs::PointCloud2 ground_msg, no_ground_msg;
  pcl::toROSMsg(*ground_cloud, ground_msg);
  pcl::toROSMsg(*filtered_no_ground_cloud, no_ground_msg);
  
  ground_msg.header = input_msg->header;
  no_ground_msg.header = input_msg->header;
  
  ground_cloud_pub_.publish(ground_msg);
  no_ground_cloud_pub_.publish(no_ground_msg);

  // Debug output
  if (debug_) {
    geometry_msgs::PoseArray debug_pose_array;
    debug_pose_array.header = input_msg->header;
    debug_pose_array.poses.push_back(getDebugPose(plane_affine));
    
    publishDebugMessage(debug_pose_array, *ground_cloud, input_msg->header);
  }
}

void RANSACGroundFilter::extractPointsIndices(
  const pcl::PointCloud<PointType>::Ptr in_cloud_ptr,
  const pcl::PointIndices & in_indices,
  pcl::PointCloud<PointType>::Ptr out_only_indices_cloud_ptr,
  pcl::PointCloud<PointType>::Ptr out_removed_indices_cloud_ptr)
{
  extract_indices_.setInputCloud(in_cloud_ptr);
  extract_indices_.setIndices(boost::make_shared<pcl::PointIndices>(in_indices));
  
  // Extract ground points
  extract_indices_.setNegative(false);
  extract_indices_.filter(*out_only_indices_cloud_ptr);
  
  // Extract non-ground points
  extract_indices_.setNegative(true);
  extract_indices_.filter(*out_removed_indices_cloud_ptr);
}

Eigen::Affine3d RANSACGroundFilter::getPlaneAffine(
  const pcl::PointCloud<PointType> segment_ground_cloud,
  const Eigen::Vector3d & plane_normal)
{
  // 修改为与ROS2版本一致的手动质心计算方法
  pcl::CentroidPoint<pcl::PointXYZ> centroid;
  for (const auto p : segment_ground_cloud.points) {
    centroid.add(p);
  }
  pcl::PointXYZ centroid_point;
  centroid.get(centroid_point);
  Eigen::Translation<double, 3> trans(centroid_point.x, centroid_point.y, centroid_point.z);
  
  PlaneBasis basis = getPlaneBasis(plane_normal);
  
  Eigen::Matrix3d rot;
  rot << basis.e_x.x(), basis.e_y.x(), basis.e_z.x(), 
          basis.e_x.y(), basis.e_y.y(), basis.e_z.y(),
          basis.e_x.z(), basis.e_y.z(), basis.e_z.z();
  
  return trans * rot;
}

void RANSACGroundFilter::applyRANSAC(
  const pcl::PointCloud<PointType>::Ptr & input,
  pcl::PointIndices::Ptr & output_inliers,
  pcl::ModelCoefficients::Ptr & output_coefficients)
{
  seg_.setInputCloud(input);
  seg_.segment(*output_inliers, *output_coefficients);
}

void RANSACGroundFilter::publishDebugMessage(
  const geometry_msgs::PoseArray & debug_pose_array,
  const pcl::PointCloud<PointType> & ground_cloud,
  const std_msgs::Header & header)
{
  debug_pose_array_pub_.publish(debug_pose_array);
  
  sensor_msgs::PointCloud2 debug_ground_msg;
  pcl::toROSMsg(ground_cloud, debug_ground_msg);
  debug_ground_msg.header = header;
  debug_ground_cloud_pub_.publish(debug_ground_msg);
}

} // namespace autoware_ground_segmentation

int main(int argc, char** argv)
{
  ros::init(argc, argv, "ransac_ground_filter_node");
  
  ros::NodeHandle nh;
  ros::NodeHandle private_nh("~");
  
  autoware_ground_segmentation::RANSACGroundFilter filter(nh, private_nh);
  
  ros::spin();
  
  return 0;
} 